//! Nuclei扫描引擎插件
//! 提供基于Nuclei模板的漏洞扫描功能

use anyhow::{anyhow, Result};
use log::{debug, error, info, warn};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Semaphore;
use tokio::time::Instant;

use super::nuclei_core::{NucleiCoreEngine, NucleiRequest, NucleiTemplate, ResponseInfo};

/// 扫描结果结构定义
#[derive(Debug, Clone, Serialize)]
pub struct ScanResult {
    /// 目标URL
    pub target: String,
    /// 漏洞ID
    pub vulnerability_id: String,
    /// 漏洞名称
    pub vulnerability_name: String,
    /// 严重程度
    pub severity: String,
    /// 匹配的URL
    pub matched_url: String,
    /// 发现时间
    pub discovered_at: String,
}

/// Nuclei扫描引擎插件结构定义
#[derive(Debug, Clone)]
pub struct NucleiEnginePlugin {
    /// 插件名称
    pub name: String,
    /// 支持的端口列表
    pub ports: Vec<u16>,
    /// 支持的扫描模式
    pub modes: Vec<String>,
    /// 模板目录路径
    pub template_dir: String,
    /// 超时时间（秒）
    pub timeout_seconds: u64,
    /// 并发工作线程数
    pub workers: usize,
    /// HTTP客户端
    pub client: reqwest::Client,
    /// Nuclei核心引擎
    pub core_engine: NucleiCoreEngine,
}

impl NucleiEnginePlugin {
    /// 创建新的Nuclei扫描引擎插件实例
    /// 
    /// # Returns
    /// - 配置默认值的NucleiEnginePlugin实例
    pub fn new() -> Result<Self> {
        // 使用项目根目录的pocs文件夹，这样更新POC文件不需要重新编译
        let default_template_dir = "./pocs";
        
        // 检查模板目录是否存在，如果不存在则创建
        if !Path::new(default_template_dir).exists() {
            match std::fs::create_dir_all(default_template_dir) {
                Ok(_) => info!("Created default template directory: {}", default_template_dir),
                Err(e) => warn!("Failed to create default template directory: {}", e),
            }
        }
        
        Ok(Self {
            name: "nuclei".to_string(),
            ports: vec![80, 443, 8080, 8443],
            modes: vec!["http", "https", "web"].iter().map(|s| s.to_string()).collect(),
            template_dir: default_template_dir.to_string(),
            timeout_seconds: 10,
            workers: 10,
            client: reqwest::Client::builder()
                .timeout(Duration::from_secs(10))
                .build()?,
            core_engine: NucleiCoreEngine::new(),
        })
    }

    /// 配置模板目录
    /// 
    /// # Arguments
    /// - `template_dir`: 模板目录路径
    /// 
    /// # Returns
    /// - 更新后的插件实例
    pub fn with_template_dir(mut self, template_dir: &str) -> Result<Self> {
        if !Path::new(template_dir).exists() || !Path::new(template_dir).is_dir() {
            return Err(anyhow!("Template directory does not exist or is not a directory: {}", template_dir));
        }
        
        self.template_dir = template_dir.to_string();
        Ok(self)
    }

    /// 配置超时时间
    /// 
    /// # Arguments
    /// - `timeout_seconds`: 超时时间，单位为秒
    /// 
    /// # Returns
    /// - 更新后的插件实例
    pub fn with_timeout(mut self, timeout_seconds: u64) -> Self {
        self.timeout_seconds = timeout_seconds;
        self.client = reqwest::Client::builder()
            .timeout(Duration::from_secs(timeout_seconds))
            .build()
            .expect("Failed to build HTTP client");
        self.core_engine = NucleiCoreEngine::with_timeout(timeout_seconds);
        self
    }

    /// 配置并发工作线程数
    /// 
    /// # Arguments
    /// - `workers`: 工作线程数
    /// 
    /// # Returns
    /// - 更新后的插件实例
    pub fn with_workers(mut self, workers: usize) -> Self {
        self.workers = workers;
        self
    }

    /// 执行扫描
    /// 
    /// # Arguments
    /// - `target`: 目标URL
    /// 
    /// # Returns
    /// - 扫描结果列表
    pub async fn scan(&mut self, target: &str) -> Result<Vec<ScanResult>> {
        info!("Starting Nuclei scan against target: {}", target);
        let start_time = Instant::now();
        
        // 加载模板
        let templates = match self.core_engine.load_templates_from_dir(&self.template_dir) {
            Ok(templates) => {
                info!("Loaded {} templates from {}", templates.len(), self.template_dir);
                templates
            },
            Err(e) => {
                error!("Failed to load templates: {}", e);
                return Err(anyhow!("Failed to load templates: {}", e));
            }
        };
        
        if templates.is_empty() {
            warn!("No templates found in {}", self.template_dir);
            return Ok(Vec::new());
        }
        
        // 准备变量
        let variables = self.core_engine.prepare_variables(target);
        
        // 使用信号量控制并发
        let semaphore = Arc::new(Semaphore::new(self.workers));
        let mut tasks = Vec::new();
        
        // 为每个模板创建一个异步任务
        for (template_path, template) in templates {
            let target = target.to_string();
            let variables = variables.clone();
            let client = self.client.clone();
            let core_engine = self.core_engine.clone();
            let semaphore = semaphore.clone();
            
            let task = tokio::spawn(async move {
                let _permit = match semaphore.acquire().await {
                    Ok(permit) => permit,
                    Err(e) => {
                        error!("Failed to acquire semaphore: {}", e);
                        return Vec::new();
                    }
                };
                
                match scan_template(&target, &template, &variables, &client, core_engine).await {
                    Ok(results) => results,
                    Err(e) => {
                        warn!("Failed to scan template {}: {}", template_path, e);
                        Vec::new()
                    }
                }
            });
            
            tasks.push(task);
        }
        
        // 等待所有任务完成并收集结果
        let mut results = Vec::new();
        for task in tasks {
            if let Ok(mut task_results) = task.await {
                results.append(&mut task_results);
            }
        }
        
        let duration = start_time.elapsed();
        info!("Nuclei scan completed in {:?}, found {} vulnerabilities", duration, results.len());
        
        Ok(results)
    }
}

/// 扫描单个模板
async fn scan_template(
    target: &str,
    template: &NucleiTemplate,
    variables: &HashMap<String, String>,
    client: &reqwest::Client,
    mut core_engine: NucleiCoreEngine,
) -> Result<Vec<ScanResult>> {
    debug!("Scanning template: {}", template.info.id);
    
    // 准备请求
    let request = prepare_request(template, variables, &mut core_engine, target)?;
    
    // 发送请求
    let response_info = match send_request(client, &request, target).await {
        Ok(info) => info,
        Err(e) => {
            debug!("Failed to send request for template {}: {}", template.info.id, e);
            return Ok(Vec::new());
        }
    };
    
    // 执行匹配器
    if core_engine.execute_matchers(&template.matchers, &response_info) {
        // 创建扫描结果
        let result = ScanResult {
            target: target.to_string(),
            vulnerability_id: template.info.id.clone(),
            vulnerability_name: template.info.name.clone(),
            severity: template.info.severity.clone(),
            matched_url: response_info.url,
            discovered_at: chrono::Utc::now().to_rfc3339(),
        };
        
        info!("Vulnerability found: {} ({} - {})", 
            template.info.name,
            template.info.id,
            template.info.severity
        );
        
        Ok(vec![result])
    } else {
        Ok(Vec::new())
    }
}

/// 准备HTTP请求
fn prepare_request(
    template: &NucleiTemplate,
    variables: &HashMap<String, String>,
    core_engine: &mut NucleiCoreEngine,
    target: &str,
) -> Result<NucleiRequest> {
    let mut request = template.request.clone();
    
    // 替换路径中的变量
    request.path = core_engine.replace_variables(&request.path, variables);
    
    // 替换头部中的变量
    let mut processed_headers = HashMap::new();
    for (key, value) in request.headers.iter() {
        processed_headers.insert(
            key.to_string(),
            core_engine.replace_variables(value, variables)
        );
    }
    request.headers = processed_headers;
    
    // 替换请求体中的变量
    request.body = core_engine.replace_variables(&request.body, variables);
    
    // 确保method有默认值
    if request.method.is_empty() {
        request.method = "GET".to_string();
    }
    
    Ok(request)
}

/// 发送HTTP请求
async fn send_request(
    client: &reqwest::Client,
    request: &NucleiRequest,
    target: &str,
) -> Result<ResponseInfo> {
    let start_time = Instant::now();
    
    // 构建完整URL
    let full_url = if request.path.starts_with("http://") || request.path.starts_with("https://") {
        request.path.clone()
    } else {
        let mut url = target.to_string();
        if !url.ends_with('/') && !request.path.starts_with('/') {
            url.push('/');
        } else if url.ends_with('/') && request.path.starts_with('/') {
            url.pop();
        }
        url + &request.path
    };
    
    debug!("Sending {} request to {}", request.method, full_url);
    
    // 构建请求构建器
    let mut builder = client.request(
        reqwest::Method::from_bytes(request.method.as_bytes())?,
        &full_url
    );
    
    // 添加请求头
    for (key, value) in &request.headers {
        builder = builder.header(key, value);
    }
    
    // 添加请求体
    if !request.body.is_empty() {
        builder = builder.body(request.body.clone());
    }
    
    // 发送请求
    let response = match builder.send().await {
        Ok(response) => response,
        Err(e) => {
            return Err(anyhow!("Failed to send request: {}", e));
        }
    };
    
    // 获取响应信息
    let status = response.status().as_u16();
    let headers: HashMap<String, String> = response.headers()
        .iter()
        .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string()))
        .collect();
    
    let body = match response.text().await {
        Ok(body) => body,
        Err(e) => {
            warn!("Failed to read response body: {}", e);
            String::new()
        }
    };
    
    let response_time = start_time.elapsed().as_millis() as u64;
    
    Ok(ResponseInfo {
        status,
        headers,
        body,
        url: full_url,
        response_time,
    })
}

// 添加缺失的导入
use serde::Serialize;
use chrono;
